﻿using LibEBML.Common;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.Windows.Media;

namespace TestPhone
{
    public class MyMediaStreamSource : MediaStreamSource
    {
        protected override void CloseMedia()
        {
            document = null;
        }

        protected override void GetDiagnosticAsync(MediaStreamSourceDiagnosticKind diagnosticKind)
        {
            throw new NotImplementedException();
        }

        protected byte[] GetNALU(int NALULength, byte[] Buff, int offset)
        {
            byte[] b = new byte[NALULength];
            Buffer.BlockCopy(Buff, offset, b, 0, b.Length);
            var len2 = ElementFactory.ConvertVintToLong(b);
            b = new byte[len2];
            Buffer.BlockCopy(Buff, offset + NALULength, b, 0, b.Length);
            return b;
            //return len2;
        }
        long FramePos;
        private void FillVideoQueue()
        {
            SimpleBlock curblock = null;
            long lp;
            while (true)
            {
                if (Clusters.Count > ClutserIndex)
                {
                    var c = Clusters[ClutserIndex];
                    if (c.Blocks.Count > BlockIndex)
                    {
                        var b = c.Blocks[BlockIndex];
                        BlockIndex++;
                        if (b.TrackNumber == VideoTrack.TrackNumber)
                        {
                            curblock = b;
                            lp = c.TimeCode + curblock.TimeCode;
                            break;
                        }
                    }
                    else
                    {
                        ClutserIndex++;
                        BlockIndex = 0;
                    }
                }
                else
                {
                    VideoQueue.Enqueue(new FrameData() { });
                    return;// null;
                }
            }
            System.IO.MemoryStream ms = new System.IO.MemoryStream();
            foreach (var item1 in curblock.GetBlockData(CompressHeader))
            {
                int i2 = 0, i = item1.Length;

                while (i2 < i)
                {
                    byte[] nalu = GetNALU(NALULength, item1, i2);
                    ms.Write(NALUSpliter, 0, 3);
                    ms.Write(nalu, 0, nalu.Length);
                    i2 += (nalu.Length + NALULength);
                }
            }
            ms.Position = 0;
            VideoQueue.Enqueue(new FrameData() { TimeCode = lp, Position = 0, Length = ms.Length, Data = ms });
        }
        Queue<FrameData> VideoQueue = new Queue<FrameData>();
        private void FillAudioQueue()
        {
            SimpleBlock curblock = null;
            var lp = AudioPosition;
            while (true)
            {
                if (Clusters.Count > AClutserIndex)
                {
                    var c = Clusters[AClutserIndex];
                    if (c.Blocks.Count > ABlockIndex)
                    {
                        var b = c.Blocks[ABlockIndex];
                        ABlockIndex++;
                        if (b.TrackNumber == AudioTrack.TrackNumber)
                        {
                            curblock = b;
                            lp = c.TimeCode + curblock.TimeCode;
                            break;
                        }
                    }
                    else
                    {
                        AClutserIndex++;
                        ABlockIndex = 0;
                    }
                }
                else
                {
                    return;
                }
            }

            if (curblock != null)
            {
                int i = 0;
                foreach (var item1 in curblock.GetBlockData(CompressHeader))
                {
                    System.IO.MemoryStream ms = new System.IO.MemoryStream(item1);
                    AudioQueue.Enqueue(new FrameData() { TimeCode = lp, Data = ms, Position = 0, Length = ms.Length });
                    i++;
                }
                FramePos = (AudioPosition - lp) / i;
            }
            else
            {
                AudioQueue.Enqueue(new FrameData() { });
                throw new Exception();
            }
        }
        Queue<FrameData> AudioQueue = new Queue<FrameData>();
        private void GetAudioData()
        {
            while (AudioQueue.Count <= 10)
            {
                try
                {
                    FillAudioQueue();
                }
                catch (Exception)
                {
                    return; //null;
                }

            }

        }

        private void GetVideoData()
        {
            while (VideoQueue.Count <= 10)
            {
                try
                {
                    FillVideoQueue();
                }
                catch (Exception)
                {
                    return; //null;
                }

            }
        }
        byte[] NALUSpliter = new byte[] { 0, 0, 1 };
        byte[] CompressHeader = null;
        protected override void GetSampleAsync(MediaStreamType mediaStreamType)
        {
            Dictionary<MediaSampleAttributeKeys, string> atts = new Dictionary<MediaSampleAttributeKeys, string>();
            var desc = Streams.Where(m => m.Type == mediaStreamType).First();
            MediaStreamSample sample;
            FrameData item1 = null;
            if (mediaStreamType == MediaStreamType.Video)
            {
                GetVideoData();
                item1 = VideoQueue.Dequeue();
            }
            else if (mediaStreamType == MediaStreamType.Audio)
            {
                var apos = AudioPosition;
                GetAudioData();
                item1 = AudioQueue.Dequeue();
            }
            else
            {
                return;
            }
            sample = new MediaStreamSample(desc, item1.Data, 0, item1.Data != null ? item1.Length : 0, item1.TimeCode * TimeScale, atts);
            ReportGetSampleCompleted(sample);
        }
        long Position = 0, AudioPosition = 0;
        int ClutserIndex = 0, BlockIndex = 0;
        int AClutserIndex = 0, ABlockIndex = 0;
        //bool FirstFrame = true;
        IDictionary<MediaSourceAttributesKeys, string> SourceAttributes;
        IEnumerable<MediaStreamDescription> Streams;
        protected override void OpenMediaAsync()
        {
            if (ReadSourceAttributes() && ReadStreamDesc())
            {
                BuildClusters();
                ReportOpenMediaCompleted(SourceAttributes, Streams);
            }
        }
        private void BuildClusters()
        {
            Clusters = MatroskaData.Segment.Clusters;
        }
        IList<Cluster> Clusters = new List<Cluster>();
        int NALULength = 0;
        TrackEntry VideoTrack, AudioTrack;
        private bool ReadStreamDesc()
        {
            List<MediaStreamDescription> sd = new List<MediaStreamDescription>();

            var tracks = MatroskaData.Segment.Tracks;//document.Segments[0][0x1654ae6b].FirstOrDefault() as LibEBML.Common.EBMLContainer;
            if (tracks != null)
            {

                foreach (var entry in tracks)
                {
                    Dictionary<MediaStreamAttributeKeys, string> at = new Dictionary<MediaStreamAttributeKeys, string>();
                    MediaStreamType tp = MediaStreamType.Audio;
                    var TrackType = entry.TrackType;//[0x83].FirstOrDefault() as EBMLType<ulong>;
                    switch (TrackType)
                    {
                        case 1:
                            VideoTrack = entry;
                            tp = MediaStreamType.Video;
                            at[MediaStreamAttributeKeys.VideoFourCC] = "H264";
                            if (entry.Video != null)
                            {
                                at[MediaStreamAttributeKeys.Width] = entry.Video.PixelWidth.ToString();
                                at[MediaStreamAttributeKeys.Height] = entry.Video.PixelHeight.ToString();
                            }
                            byte[] b = entry.CodecPrivate;
                            if (b[0] != 0x01)
                            {
                                ErrorOccurred("错误的CodecPrivate格式");
                                return false;
                            }
                            int mp = 5;
                            byte x1f = 0x1f;
                            NALULength = b[4] & 3 + 1;
                            var nsps = b[5] & x1f;
                            mp++;
                            var lsps = b[6] * 256 + b[7];
                            mp += 2;
                            int ml = nsps * lsps;
                            mp += ml;
                            var npps = b[mp];
                            var lpps = b[mp + 1] * 256 + b[mp + 2];
                            int ml2 = npps * lpps;
                            byte[] cpdata = new byte[ml + ml2 + 6];
                            cpdata[2] = 1;
                            Buffer.BlockCopy(b, 8, cpdata, 3, ml);
                            cpdata[ml + 5] = 1;
                            Buffer.BlockCopy(b, mp + 3, cpdata, ml + 6, ml2);
                            at[MediaStreamAttributeKeys.CodecPrivateData] = BitConverter.ToString(cpdata).Replace("-", "");

                            break;
                        case 2:
                            AudioTrack = entry;
                            tp = MediaStreamType.Audio;
                            using (var ms = new System.IO.MemoryStream())
                            {
                                var w = new System.IO.BinaryWriter(ms);
                                w.Write((short)0xff);
                                w.Write((short)entry.Audio.Channels);
                                w.Write((int)entry.Audio.SamplingFrequency);
                                w.Write((int)(entry.Audio.SamplingFrequency * entry.Audio.Channels * 0x02 / 0x04));
                                //w.Write((int)12000);
                                w.Write((short)0x04);
                                w.Write((short)0x10);
                                w.Write((short)entry.CodecPrivate.Length);
                                w.Write(entry.CodecPrivate);
                                byte[] hd = ms.ToArray();
                                at[MediaStreamAttributeKeys.CodecPrivateData] = BitConverter.ToString(hd).Replace("-", "");
                            }

                            break;
                        default:
                            //return false;
                            break;
                    }
                    MediaStreamDescription desc = new MediaStreamDescription(tp, at);
                    sd.Add(desc);
                    //    if (TrackType != null && TrackType == 1)
                    //    {
                    //        var TrackNumber = entry[0xd7].FirstOrDefault() as EBMLType<ulong>;
                    //        if (TrackNumber != null)
                    //        {
                    //            VideoTrackNumber = (long)TrackNumber.Value;
                    //        }
                    //        var ContentEncodings = entry[0x6d80].FirstOrDefault() as EBMLContainer;
                    //        if (ContentEncodings != null)
                    //        {
                    //            foreach (EBMLContainer ContentEncoding in ContentEncodings)
                    //            {
                    //                var cc = ContentEncoding[0x5034].FirstOrDefault() as EBMLContainer;
                    //                if (cc != null)
                    //                {
                    //                    var cca = cc[0x4254].FirstOrDefault() as EBMLType<ulong>;
                    //                    if (cca != null && cca.Value == 3)
                    //                    {
                    //                        var hs = cc[0x4255].FirstOrDefault() as EBMLBinary;
                    //                        CompressHeader = new byte[hs.Size];
                    //                        hs.BaseSource.Position = hs.DataPosition;
                    //                        hs.BaseSource.Read(CompressHeader, 0, CompressHeader.Length);
                    //                    }
                    //                }
                    //            }
                    //        }
                }
            }

            Streams = sd;
            return true;
        }
        long TimeScale = 1;
        private bool ReadSourceAttributes()
        {
            Dictionary<MediaSourceAttributesKeys, string> mediaStreamAttributes = new Dictionary<MediaSourceAttributesKeys, string>();
            mediaStreamAttributes.Clear();

            var d1 = MatroskaData.Segment.Info.Duration;
            var d2 = MatroskaData.Segment.Info.TimecodeScale;
            var d = d1 * d2 * 10;
            //var d = d1 * 100;
            TimeScale = (long)d2 / 100;
            mediaStreamAttributes[MediaSourceAttributesKeys.CanSeek] = "true";
            mediaStreamAttributes[MediaSourceAttributesKeys.Duration] = d.ToString("F0");

            SourceAttributes = mediaStreamAttributes;

            return true;
            //}
            //else
            //{
            //    ErrorOccurred("无效的Info信息");
            //    return false;
            //}
        }
        LibEBML.Common.EBMLDocument document = null;
        LibEBML.Common.Matroska MatroskaData = null;
        public MyMediaStreamSource(System.IO.Stream source)
        {
            LibEBML.Common.EBMLDocument doc = new LibEBML.Common.EBMLDocument(source);
            document = doc;
            MatroskaData = new Matroska(doc);
        }
        protected override void SeekAsync(long seekToTime)
        {
            var c = Clusters.Where(m => m.TimeCode * TimeScale <= seekToTime).LastOrDefault();
            if (c == null)
            {
                Position = 0;
                ClutserIndex = Clusters.Count - 1;
                BlockIndex = 0;
            }
            else
            {
                Position = c.TimeCode;
                ClutserIndex = Clusters.IndexOf(c);
                BlockIndex = 0;
            }
            AudioPosition = Position;
            AClutserIndex = ClutserIndex;
            ABlockIndex = BlockIndex;
            ReportSeekCompleted(Position * TimeScale);
        }

        protected override void SwitchMediaStreamAsync(MediaStreamDescription mediaStreamDescription)
        {
            throw new NotImplementedException();
        }
    }
}
